/*
 * Copyright (C) 2025 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use crate::common::constants::*;
use crate::traits::crypto_engine::{AesGcmParam, AesGcmResult, CryptoEngineRegistry};
use crate::traits::misc_manager::MiscManagerRegistry;
use crate::traits::time_keeper::TimeKeeperRegistry;
use crate::types::{Udid, CredentialId};
use crate::Vec;

pub const AES_GCM_TOKEN_AAD_BYTES: &[u8] = b"OH_authToken";
pub const AUTH_TOKEN_CIPHER_LEN: usize = core::mem::size_of::<TokenDataToEncrypt>();
pub const TOKEN_VERSION: u32 = 0;

//加密+签名后的authtoken产物
#[repr(C, packed)]
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct UserAuthTokenHal {
    pub version: u32,
    pub token_data_plain: TokenDataPlain,
    pub token_data_cipher: [u8; AUTH_TOKEN_CIPHER_LEN],
    pub tag: [u8; AES_GCM_TAG_SIZE],
    pub iv: [u8; AES_GCM_IV_SIZE],
    pub sign: [u8; SHA256_DIGEST_SIZE],
}

impl UserAuthTokenHal {
    pub fn new(
        version: u32,
        token_data_plain: TokenDataPlain,
        token_data_cipher: [u8; AUTH_TOKEN_CIPHER_LEN],
        tag: [u8; AES_GCM_TAG_SIZE],
        iv: [u8; AES_GCM_IV_SIZE],
        sign: [u8; SHA256_DIGEST_SIZE],
    ) -> Self {
        Self {
            version,
            token_data_plain,
            token_data_cipher,
            tag,
            iv,
            sign,
        }
    }

    //结构体序列化
    fn serialize(&self) -> &[u8] {
        unsafe {
            core::slice::from_raw_parts(
                self as *const Self as *const u8,
                core::mem::size_of::<Self>(),
            )
        }
    }

    //结构体反序列化
    fn deserialize(bytes: &[u8]) -> Result<Self, ErrorCode> {
        let expected_len = core::mem::size_of::<Self>();
        if bytes.len() != expected_len {
            return Err(ErrorCode::GeneralError);
        }

        Ok(unsafe { core::ptr::read_unaligned(bytes.as_ptr() as *const Self) })
    }

    // 对令牌进行签名并返回已签名的令牌
    fn signature(&self) -> Result<Self, ErrorCode> {
        let token_bytes = self.serialize();
        let mac_key = MiscManagerRegistry::get().get_token_key()?.mac_key;
        let sign_slice = CryptoEngineRegistry::get().hmac_sha256(&mac_key, token_bytes)?;

        if sign_slice.len() != SHA256_DIGEST_SIZE {
            return Err(ErrorCode::GeneralError);
        }

        let mut signed_token = self.clone();
        signed_token.sign.copy_from_slice(&sign_slice);

        Ok(signed_token)
    }

    // 验证令牌签名
    fn verify_signature(&self) -> Result<(), ErrorCode> {
        let token_bytes = self.serialize();
        let mac_key = MiscManagerRegistry::get().get_token_key()?.mac_key;
        let computed_signature = CryptoEngineRegistry::get().hmac_sha256(&mac_key, token_bytes)?;

        // 比较签名
        if computed_signature.as_slice() != self.sign.as_slice() {
            return Err(ErrorCode::BadSign);
        }

        Ok(())
    }
}

//authtoken中的明文
#[repr(C, packed)]
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct TokenDataPlain {
    challenge: [u8; CHALLENGE_LEN],
    time: u64,
    auth_trust_level: AuthTrustLevel,
    auth_type: AuthType,
    schedule_mode: ScheduleMode,
    security_level: AuthSecurityLevel,
    token_type: TokenType,
}

impl TokenDataPlain {
    pub fn new(
        challenge: [u8; CHALLENGE_LEN],
        time: u64,
        auth_trust_level: AuthTrustLevel,
        auth_type: AuthType,
        schedule_mode: ScheduleMode,
        security_level: AuthSecurityLevel,
        token_type: TokenType,
    ) -> Self {
        Self {
            challenge: challenge,
            time: time,
            auth_trust_level: auth_trust_level,
            auth_type: auth_type,
            schedule_mode: schedule_mode,
            security_level: security_level,
            token_type: token_type,
        }
    }
}

//authtoken中的密文
#[repr(C, packed)]
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct TokenDataToEncrypt {
    user_id: i32,
    secure_uid: u64,
    enrolled_id: u64,
    credential_id: CredentialId,
    collector_udid: Udid,
    verifier_udid: Udid
}

impl TokenDataToEncrypt {
    pub fn new(
        user_id: i32,
        secure_uid: u64,
        enrolled_id: u64,
        credential_id: CredentialId,
        collector_udid: Udid,
        verifier_udid: Udid
    ) -> Self {
        Self {
            user_id: user_id,
            secure_uid: secure_uid,
            enrolled_id: enrolled_id,
            credential_id: credential_id,
            collector_udid: collector_udid,
            verifier_udid: verifier_udid,
        }
    }

    //结构体序列化
    fn serialize(&self) -> &[u8] {
        unsafe {
            core::slice::from_raw_parts(
                self as *const Self as *const u8,
                core::mem::size_of::<Self>(),
            )
        }
    }

    //结构体反序列化
    fn deserialize(bytes: &[u8]) -> Result<Self, ErrorCode> {
        let expected_len = core::mem::size_of::<Self>();
        if bytes.len() != expected_len {
            return Err(ErrorCode::GeneralError);
        }

        Ok(unsafe { core::ptr::read_unaligned(bytes.as_ptr() as *const Self) })
    }

    // 加密
    fn encrypt(
        &self,
        aes_param: &mut AesGcmParam,
    ) -> Result<([u8; AUTH_TOKEN_CIPHER_LEN], [u8; AES_GCM_TAG_SIZE]), ErrorCode> {
        let token_data_to_encrypt_serialized = self.serialize();
        let aes_gcm_result = CryptoEngineRegistry::get()
            .aes_gcm_encrypt(&token_data_to_encrypt_serialized, aes_param)?;

        let mut ciphertext = [0u8; AUTH_TOKEN_CIPHER_LEN];
        let mut tag = [0u8; AES_GCM_TAG_SIZE];

        ciphertext.copy_from_slice(&aes_gcm_result.ciphertext);
        tag.copy_from_slice(&aes_gcm_result.authentication_tag);

        Ok((ciphertext, tag))
    }

    // 解密令牌数据
    fn decrypt(token: &UserAuthTokenHal) -> Result<Self, ErrorCode> {
        let mut aes_param = AuthTokenHelper::get_aes_gcm_param()?;
        aes_param.iv = token.iv;

        let ciphertext = token.token_data_cipher.to_vec();
        let aes_gcm_result = AesGcmResult::new(ciphertext, token.tag);

        let decrypted_data =
            CryptoEngineRegistry::get().aes_gcm_decrypt(&aes_param, &aes_gcm_result)?;

        TokenDataToEncrypt::deserialize(&decrypted_data)
    }
}

//authtoken结构体
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct AuthToken {
    pub challenge: [u8; CHALLENGE_LEN],
    pub time: u64,
    pub auth_trust_level: AuthTrustLevel,
    pub auth_type: AuthType,
    pub schedule_mode: ScheduleMode,
    pub security_level: AuthSecurityLevel,
    pub token_type: TokenType,
    pub user_id: i32,
    pub secure_uid: u64,
    pub enrolled_id: u64,
    pub credential_id: CredentialId,
    pub collector_udid: Udid,
    pub verifier_udid: Udid
}

impl AuthToken {
    pub fn new(
        challenge: [u8; CHALLENGE_LEN],
        time: u64,
        auth_trust_level: AuthTrustLevel,
        auth_type: AuthType,
        schedule_mode: ScheduleMode,
        security_level: AuthSecurityLevel,
        token_type: TokenType,
        user_id: i32,
        secure_uid: u64,
        enrolled_id: u64,
        credential_id: CredentialId,
        collector_udid: Udid,
        verifier_udid: Udid
    ) -> Self {
        Self {
            challenge,
            time,
            auth_trust_level,
            auth_type,
            schedule_mode,
            security_level,
            token_type,
            user_id,
            secure_uid,
            enrolled_id,
            credential_id,
            collector_udid,
            verifier_udid,
        }
    }

    pub fn encode_to_slice(&mut self) -> Result<Vec<u8>, ErrorCode> {
        //更新时间
        self.time = TimeKeeperRegistry::get().get_system_time()?;

        //加密
        let mut aes_param = AuthTokenHelper::get_aes_gcm_param()?;
        let token_data_to_encrypt = TokenDataToEncrypt::new(
            self.user_id,
            self.secure_uid,
            self.enrolled_id,
            self.credential_id,
            self.collector_udid,
            self.verifier_udid,
        );
        let (ciphertext, tag) = token_data_to_encrypt.encrypt(&mut aes_param)?;

        //签名
        let unsigned_token = self.create_unsigned_token(ciphertext, tag, aes_param.iv);
        let sign_auth_token = unsigned_token.signature()?;

        Ok(sign_auth_token.serialize().to_vec())
    }

    pub fn decode_from_slice(
        user_auth_token_hal_slice: &[u8],
        allowable_duration: u64,
    ) -> Result<Self, ErrorCode> {
        let user_auth_token_hal = UserAuthTokenHal::deserialize(user_auth_token_hal_slice)?;
        //检查时间是不是valid
        let is_time_valid = AuthTokenHelper::is_time_valid(
            user_auth_token_hal.token_data_plain.time,
            allowable_duration,
        )?;

        if !is_time_valid {
            return Err(ErrorCode::TokenTimeout);
        }

        //检查签名
        user_auth_token_hal.verify_signature()?;

        //解密
        let token_data_to_encrypt = TokenDataToEncrypt::decrypt(&user_auth_token_hal)?;

        Ok(Self {
            challenge: user_auth_token_hal.token_data_plain.challenge,
            time: user_auth_token_hal.token_data_plain.time,
            auth_trust_level: user_auth_token_hal.token_data_plain.auth_trust_level,
            auth_type: user_auth_token_hal.token_data_plain.auth_type,
            schedule_mode: user_auth_token_hal.token_data_plain.schedule_mode,
            security_level: user_auth_token_hal.token_data_plain.security_level,
            token_type: user_auth_token_hal.token_data_plain.token_type,
            user_id: token_data_to_encrypt.user_id,
            secure_uid: token_data_to_encrypt.secure_uid,
            enrolled_id: token_data_to_encrypt.enrolled_id,
            credential_id: token_data_to_encrypt.credential_id,
            collector_udid: token_data_to_encrypt.collector_udid,
            verifier_udid: token_data_to_encrypt.verifier_udid,
        })
    }

    // 创建未签名的令牌结构
    fn create_unsigned_token(
        &self,
        ciphertext: [u8; AUTH_TOKEN_CIPHER_LEN],
        tag: [u8; AES_GCM_TAG_SIZE],
        iv: [u8; AES_GCM_IV_SIZE],
    ) -> UserAuthTokenHal {
        let token_data_plain = TokenDataPlain::new(
            self.challenge,
            self.time,
            self.auth_trust_level,
            self.auth_type,
            self.schedule_mode,
            self.security_level,
            self.token_type,
        );
        UserAuthTokenHal::new(
            TOKEN_VERSION,
            token_data_plain,
            ciphertext,
            tag,
            iv,
            [0; SHA256_DIGEST_SIZE],
        )
    }

    //结构体序列化
    fn serialize(&self) -> &[u8] {
        unsafe {
            core::slice::from_raw_parts(
                self as *const Self as *const u8,
                core::mem::size_of::<Self>(),
            )
        }
    }

    //结构体反序列化
    fn deserialize(bytes: &[u8]) -> Result<Self, ErrorCode> {
        let expected_len = core::mem::size_of::<Self>();
        if bytes.len() != expected_len {
            return Err(ErrorCode::GeneralError);
        }

        Ok(unsafe { core::ptr::read_unaligned(bytes.as_ptr() as *const Self) })
    }
}

struct AuthTokenHelper;

impl AuthTokenHelper {
    //初始化加密参数
    fn get_aes_gcm_param() -> Result<AesGcmParam, ErrorCode> {
        let token_key = MiscManagerRegistry::get().get_token_key()?;

        Ok(AesGcmParam {
            key: token_key.cipher_key.to_vec().into(),
            iv: [0u8; AES_GCM_IV_SIZE],
            aad: AES_GCM_TOKEN_AAD_BYTES.to_vec().into(),
        })
    }

    //检查时间是否有效
    fn is_time_valid(time: u64, allowable_duration: u64) -> Result<bool, ErrorCode> {
        let current_time = TimeKeeperRegistry::get().get_system_time()?;

        if current_time < time {
            return Ok(false);
        }

        if current_time - time > allowable_duration {
            return Ok(false);
        }

        Ok(true)
    }
}
